import os
import sys
import warnings
warnings.filterwarnings("ignore")

import numpy as np # Numerical computing
import pandas as pd # Data manipulatio

import seaborn as sns # R-like high-level plots
import matplotlib.pyplot as plt # Python's base plotting 

sys.path.append('../')

# Utility imports for this analysis
from src import utils, plots

# Initialize the timer
startTime = utils.getTime()

## Figure Settings

# Define default colors and styles for plots
def_colors = [
    "#139593", "#fca311", "#e54f2a",
    "#c3c3c3", "#555555",
    "#690000", "#5f4a00", "#004549"
]

# Set seaborn style
sns.set_theme(
    style="white",
    context="paper",
    palette=def_colors,
    font_scale=1,
    rc={
        "figure.figsize": (6, 4),
        "font.family": "sans-serif",
        "font.sans-serif": ["Arial", "Ubuntu Mono"],
    }
)

# Figure Saving Settings
figure_formats = ["pdf"]
save_to_folder = True
transparent_bg = True
figure_dpi = 300

## Configure dataframe displaying
pd.set_option('display.float_format', lambda x: '%.4f' % x)
pd.set_option('display.max_rows', 50)
pd.set_option('display.max_columns', 50)
pd.set_option('display.width', 500)  # Set a wider display width

## Printing Settings
verbose = True

## Show the color-palettes
plots.color_palette( def_colors, save=False )

# Establish paths and 
notebook_name = "05-IdBenchmark"
input_path = f"./data/"
mainFig_path = f"./figures/"

def setup_simulation_paths( simID ):
    """
    Create output and figure directories for a simulation, if they do not exist. 
    (Uses global variable for save_to_folder and figure_formats)

    Args:
        simID (str): Simulation ID.
        
        (Global Variables)
            input_path (str): Base path for input data.
            mainFig_path (str): Base path for figures.
            save_to_folder (bool): Whether to save figures to folders.
            figure_formats (list): List of formats to save figures in.
    
    Returns:
        tuple: (output_path, figure_path)
    """
    output_path = f"{input_path}{simID}/"
    if not os.path.exists(output_path):
        os.makedirs(output_path)

    figure_path = f"{mainFig_path}{simID}/"
    if save_to_folder:
        for fmt in figure_formats:
            cur_folder = os.path.join(figure_path, fmt)
            if not os.path.exists(cur_folder):
                os.makedirs(cur_folder)
    return output_path, figure_path

# Global variables
seed = 42 # Seed for reproducibility  
pthr = 10**-3  # p-value threshold for significance
thresholds = list(utils.generate_thresholds(10.0, -15, 1, 0, 1, 0.1)) # Thresholds for the analysis

# Methods and their plotting styles
method_palette = {
    "COPF": "#139593",
    "PeCorA": "#fca311",
    "ProteoForge": "#e54f2a",
}
method_styles = {
    "COPF": "--",
    "PeCorA": "-.",
    "ProteoForge": ":",
}
method_markers = {
    "COPF": "o",
    "PeCorA": "s",
    "ProteoForge": "^",
}


# Matthews Correlation Coefficient (MCC) thresholds and colors
mcc_thresholds = {
    0.0 : 'Random',
    0.3 : 'Fair (0.3)',
    0.5 : 'Moderate (0.5)',
    0.7 : 'Strong (0.7)',
    0.9 : 'Almost Perfect (0.9)',
}
mcc_colors = {
    'Random': '#c3c3c3',             
    'Fair (0.3)': '#e0aaff',         
    'Moderate (0.5)': '#9d4edd',     
    'Strong (0.7)': '#5a189a',       
    'Almost Perfect (0.9)': '#240046' 
}

plots.color_palette(mcc_colors, save=False, name="MCC Interpretation Colors")

stTime = utils.getTime()

simID = "Sim1"  
# Set up a path for the simulation
output_path, figure_path = setup_simulation_paths( simID )

methods = list(method_palette.keys())
dataTypes = ['complete', 'imputed']
experiments = ['twoPep', 'randomPep', 'halfPep', 'halfPlusPep']
experiment_mapper = {
    'twoPep': 'Two Peptides', 
    'randomPep': '2>50% Peptides',
    'halfPep': '50% Peptides',
    'halfPlusPep': '>50% Peptides',
}

# =============================================================================
# ANALYSIS SETUP
# =============================================================================
print("PEPTIDE IDENTIFICATION BENCHMARK ANALYSIS")
print("-" * 52)
print(f"Simulation ID: {simID}")
print(f"Methods: {', '.join(methods)} ({len(methods)} total)")
print(f"Data Types: {', '.join(dataTypes)}")
print(f"Experiments: {len(experiments)} scenarios")
print(f"Output Path: {output_path}")

# Progress tracking
total_combinations = len(methods) * len(dataTypes) * len(experiments)
current_combination = 0
files_found = 0
files_missing = 0

print(f"Total combinations to process: {total_combinations}")
print("-" * 52)
# =============================================================================
# DATA PROCESSING
# =============================================================================
display("Calculating ID Benchmarks:")

combined_results = []
for cur_method in methods:
    for cur_type in dataTypes:
        for cur_exp in experiments:
            current_combination += 1
            
            # Load the result data
            result_file = f"{output_path}2_{cur_exp}_{cur_type}_{cur_method}_ResultData.feather"
            
            if os.path.exists(result_file):
                res_df = pd.read_feather(result_file)
                files_found += 1
                
                # Show progress for successful loads only
                if verbose:
                    progress = f"[{current_combination:2d}/{total_combinations}]"
                    print(f"{progress} ✓ {cur_method:12} | {cur_type:9} | {experiment_mapper[cur_exp]}")
            else:
                files_missing += 1
                if verbose:
                    progress = f"[{current_combination:2d}/{total_combinations}]"
                    print(f"{progress} ✗ {cur_method:12} | {cur_type:9} | {experiment_mapper[cur_exp]} (MISSING)")
                continue
            
            # Handle COPF-specific column renaming
            if cur_method == "COPF":
                res_df = res_df.rename(columns={
                    "proteoform_score_pval": "adj_pval",
                    'protein_id': "Protein"
                })

            # Create metric data
            metric_data = utils.create_metric_data(
                res_df, 
                pvalue_thresholds=thresholds,
                label_col='pertPeptide',
                pvalue_col='adj_pval'
            )
            metric_data['Method'] = cur_method
            metric_data['DataType'] = cur_type
            metric_data['Experiment'] = cur_exp
            combined_results.append(metric_data)

# Combine and process data
idBenchmarkData = pd.concat(combined_results, ignore_index=True)
idBenchmarkData['Experiment'] = idBenchmarkData['Experiment'].map(experiment_mapper)

# Categorical ordering ensure twoPep, randomPep, halfPep, halfPlusPep
idBenchmarkData['Experiment'] = pd.Categorical(
    idBenchmarkData['Experiment'], 
    categories=[experiment_mapper[exp] for exp in experiments],
    ordered=True
)

print(f"\nDATA PREVIEW:")
display(idBenchmarkData.tail(5))

# Save the processed data
output_file = f"{output_path}4_{simID}_PeptideIdentification_PerformanceData.feather"
idBenchmarkData.to_feather(output_file)
idBenchmarkData.to_csv(output_file.replace('.feather', '.csv'), index=False)

print(f"\nRESULTS SUMMARY:")
print("=" * 52)
execTime = utils.prettyTimer(utils.getTime() - stTime)
print(f"Total execution time:         {execTime}")
print(f"Files processed successfully: {files_found}")
print(f"Files missing/skipped:        {files_missing}")
print(f"Processing success rate:      {files_found/(files_found+files_missing)*100:.1f}%")
print("-" * 52)
print(f"Final dataset shape:          {idBenchmarkData.shape}")
print(f"Unique experiments:           {idBenchmarkData['Experiment'].nunique()}")
print(f"Methods analyzed:             {idBenchmarkData['Method'].nunique()}")
print(f"Data types included:          {idBenchmarkData['DataType'].nunique()}")
print(f"Data saved to: {output_file}")
print("=" * 52)

# ROC Curves for different perturbations and methods
fig, axes = plt.subplots(
    2, 4, figsize=(16, 6), 
    sharey=True, sharex=True,
    gridspec_kw={
        "wspace": 0.05, 
        "hspace": 0.05
    }
)

for i, pert in enumerate(idBenchmarkData['Experiment'].unique()):
    for j, dataType in enumerate(dataTypes):
        cur_data = idBenchmarkData[
            (idBenchmarkData["Experiment"] == pert) & 
            (idBenchmarkData["DataType"] == dataType)
        ]
        # Ensure the data is complete for ROC curves
        cur_data = cur_data.groupby("Method").apply(
            lambda x: utils.complete_curve_data(x, 'ROC', 'FPR', 'TPR')
        ).reset_index(drop=True)
        
        # Calculate AUC per method from TPR and FPR
        auc_data = cur_data.groupby("Method").apply(
            lambda x: np.trapezoid(
                x.sort_values("FPR")["TPR"], x.sort_values("FPR")["FPR"]
            ), 
            include_groups=False
        )
        # Plot the ROC curve
        sns.lineplot(
            data=cur_data,
            x="FPR",
            y="TPR",
            hue="Method",
            style="Method",
            palette=method_palette,
            ax=axes[j, i],
            # Show the points
            markers=method_markers,
            dashes=False,
            markersize=7.5,
            markeredgewidth=0.5,
            markeredgecolor="gray", 
            legend=False,
            rasterized=True,
            estimator=None
        )
        # Add AUC values as legend like text
        for k, method in enumerate(auc_data.index):
            auc = auc_data[method]
            # Add color to match the palette
            color = method_palette[method]
            axes[j, i].text(
                0.95,
                0.05 + k * 0.075,
                f"{method}: {auc:.2f}",
                color=color,
                transform=axes[j, i].transAxes,
                ha="right",
                va="bottom",
                fontsize=12,
                fontweight="bold",
            )
            # Add a large no-facecolor marker to the pthr value for each method
            pthr_data = cur_data[(cur_data["Method"] == method) & (cur_data["threshold"] == pthr)]
            axes[j, i].scatter(
                pthr_data["FPR"],
                pthr_data["TPR"],
                color=color,
                s=125,
                edgecolor="black",
                linewidth=1.25,
                marker="o",
                facecolors="none",
            )
        
        # Add the diagonal line
        axes[j, i].plot([0, 1], [0, 1], color="black", linestyle="--")
        # axes[j, i].set_title(f"{pert}", fontsize=14, fontweight="bold")
        axes[j, i].set_xlabel("FPR", fontsize=12)
        axes[j, i].set_ylabel("TPR", fontsize=12)
        axes[j, i].set_xlim(-0.05, 1.05)
        axes[j, i].set_ylim(-0.05, 1.05)
        axes[j, i].grid("both", linestyle="--", linewidth=0.75, alpha=0.5, color="lightgrey")
        # If top row set the title
        if j == 0:
            axes[j, i].set_title(f"{pert}", fontsize=14, fontweight="bold")
        # For the last column, set the text for the dataType on the right
        if i == 3:
            axes[j, i].text(
                1.05,
                0.5,
                f"{dataType.capitalize()}",
                transform=axes[j, i].transAxes,
                ha="left",
                va="center",
                fontsize=14,
                fontweight="bold",
                rotation=270,
            )

# Add a text indicating the circle is the p-value threshold under fig.title and above the first subplot
fig.text(
    0.9,
    .97,
    f"Black circles represent the TPR and FPR values at p-value threshold of {pthr}",
    ha="right",
    va="center",
    fontsize=12,
    fontstyle="italic",
    color="gray",
)

# Finalize the plot
plt.tight_layout()
plots.finalize_plot(
    fig, 
    show=True,
    save=save_to_folder,
    filename=f"{simID}_PepIDBenchmark_ROC_Curve", 
    filepath=figure_path,
    formats=figure_formats, 
    transparent=transparent_bg,
    dpi=figure_dpi,
)

# PR Curves for different perturbations and methods
fig, axes = plt.subplots(
    2, 4, figsize=(16, 6), 
    sharey=True, sharex=True,
    gridspec_kw={
        "wspace": 0.05, 
        "hspace": 0.05
    }
)

for i, pert in enumerate(idBenchmarkData['Experiment'].unique()):
    for j, dataType in enumerate(dataTypes):
        cur_data = idBenchmarkData[
            (idBenchmarkData["Experiment"] == pert) & 
            (idBenchmarkData["DataType"] == dataType)
        ]
        
        # Add a small value to avoid division by zero
        cur_data["Precision"] = cur_data["TP"] / (cur_data["TP"] + cur_data["FP"] + 1e-6)
        cur_data["Recall"] = cur_data["TP"] / (cur_data["TP"] + cur_data["FN"] + 1e-6)
        
        cur_data = cur_data.sort_values("Recall", ascending=False)
        cur_data = cur_data[~((cur_data["Recall"] == 0) & (cur_data["Precision"] == 0))]

        # Ensure the data is complete for PR curves
        cur_data = cur_data.groupby("Method").apply(
            lambda x: utils.complete_curve_data(x, 'PR', 'Recall', 'Precision')
        ).reset_index(drop=True)
        
        # Calculate AUC per method from Precision and Recall
        # Here we use the trapezoidal rule to calculate the area under the curve
        f1_data = cur_data.groupby("Method").apply(
            lambda x: np.trapezoid(
                x.sort_values("Recall")["Precision"], x.sort_values("Recall")["Recall"]
            ),
            include_groups=False
        )

        # Plot the PR curve
        sns.lineplot(
            data=cur_data,
            x="Recall",
            y="Precision",
            hue="Method",
            style="Method",
            palette=method_palette,
            ax=axes[j, i],
            # Show the points
            markers=method_markers,
            dashes=False,
            markersize=7.5,
            markeredgewidth=0.5,
            markeredgecolor="gray", 
            legend=False,
            rasterized=True,
            estimator=None
        )

        # Add AUC values as legend like text
        for k, method in enumerate(f1_data.index):
            f1 = f1_data[method]
            # Add color to match the palette
            color = method_palette[method]
            axes[j, i].text(
                0.05,
                0.05 + k * 0.075,
                f"{method}: {f1:.2f}",
                color=color,
                transform=axes[j, i].transAxes,
                ha="left",
                va="bottom",
                fontsize=12,
                fontweight="bold",
            )
            # Add a large no-facecolor marker to the pthr value for each method
            pthr_data = cur_data[(cur_data["Method"] == method) & (cur_data["threshold"] == pthr)]
            axes[j, i].scatter(
                pthr_data["Recall"],
                pthr_data["Precision"],
                color=color,
                s=125,
                edgecolor="black",
                linewidth=1.25,
                marker="o",
                facecolors="none",
            )

        axes[j, i].set_xlabel("Recall", fontsize=12)
        axes[j, i].set_ylabel("Precision", fontsize=12)
        axes[j, i].set_xlim(-0.05, 1.05)
        axes[j, i].set_ylim(-0.05, 1.05)
        axes[j, i].grid("both", linestyle="--", linewidth=0.75, alpha=0.5, color="lightgrey")
        # If top row set the title
        if j == 0:
            axes[j, i].set_title(f"{pert}", fontsize=14, fontweight="bold")
        # For the last column, set the text for the dataType on the right
        if i == 3:
            axes[j, i].text(
                1.05,
                0.5,
                f"{dataType.capitalize()}",
                transform=axes[j, i].transAxes,
                ha="left",
                va="center",
                fontsize=14,
                fontweight="bold",
                rotation=270,
            )

# Add a text indicating the circle is the p-value threshold under fig.title and above the first subplot
fig.text(
    0.9,
    .97,
    f"Black circles represent the Precision and Recall values at p-value threshold of {pthr}",
    ha="right",
    va="center",
    fontsize=12,
    fontstyle="italic",
    color="gray",
)
# Finalize the plot
plt.tight_layout()
plots.finalize_plot(
    fig, 
    show=True,
    save=save_to_folder,
    filename=f"{simID}_PepIDBenchmark_PR_Curve", 
    filepath=figure_path,
    formats=figure_formats, 
    transparent=transparent_bg,
    dpi=figure_dpi,
)

best_thresholds = idBenchmarkData.groupby(["Method", "Experiment", "DataType"]).apply(
    lambda x: x.loc[x["MCC"].idxmax(), ["threshold", "MCC"]]
).reset_index()
best_thresholds.columns = ["Method", "Experiment", "DataType", "threshold", "MCC"]
# print("\nBest thresholds per method and experiment:")
# display(best_thresholds)

# MCC Curves for p-value thresholds
fig, axes = plt.subplots(
    2, 4, figsize=(16, 6), 
    sharey=True, sharex=True,
    gridspec_kw={
        "wspace": 0.05, 
        "hspace": 0.05
    }
)
for i, pert in enumerate(idBenchmarkData['Experiment'].unique()):
    for j, dataType in enumerate(dataTypes):
        cur_data = idBenchmarkData[
            (idBenchmarkData["Experiment"] == pert) & 
            (idBenchmarkData["DataType"] == dataType)
        ].copy()
        cur_data['-log10(threshold)'] = -np.log10(cur_data['threshold'])
        # Plot the MCC curve
        sns.lineplot(
            data=cur_data,
            x="-log10(threshold)",
            y="MCC",
            hue="Method",
            style="Method",
            palette=method_palette,
            ax=axes[j, i],
            # Show the points
            markers=method_markers,
            dashes=False,
            markersize=7.5,
            markeredgewidth=0.5,
            markeredgecolor="gray", 
            legend=False,
            rasterized=True,
            estimator=None
        )
        # Add the best MCC point for each method
        for k, method in enumerate(cur_data["Method"].unique()):
            best_data = best_thresholds[
                (best_thresholds["Method"] == method) & 
                (best_thresholds["Experiment"] == pert) & 
                (best_thresholds["DataType"] == dataType)
            ]
            if not best_data.empty:
                best_threshold = best_data["threshold"].values[0]
                best_mcc = best_data["MCC"].values[0]
                axes[j, i].scatter(
                    -np.log10(best_threshold),
                    best_mcc,
                    color=method_palette[method],
                    s=125,
                    edgecolor="black",
                    linewidth=1.25,
                    marker="o",
                    facecolors="none",
                )
                # # Annotate the best MCC point
                # Add color to match the palette
                color = method_palette[method]
                axes[j, i].text(
                    0.95,
                    0.75 + k * 0.075,
                    f"{method}: {best_mcc:.2f}",
                    color=color,
                    ha="right",
                    va="bottom",
                    fontsize=12,
                    fontweight="bold",
                    transform=axes[j, i].transAxes,
                )

        axes[j, i].set_xlabel("-log10(p-value threshold)", fontsize=12)
        axes[j, i].set_ylabel("MCC", fontsize=12)
        axes[j, i].set_xlim(-0.5, 15.5)
        axes[j, i].set_ylim(-0.25, 1.05)
        axes[j, i].grid("both", linestyle="--", linewidth=0.75, alpha=0.5, color="lightgrey")
        # Add horizontal line at y=0
        # axes[j, i].axhline(0, color="black", linestyle="--", linewidth=1)
        # Draw MCC interpretation thresholds
        for thresh, label in mcc_thresholds.items():
            axes[j, i].axhline(
                thresh, color=mcc_colors[label], alpha=1,
                linestyle="dotted", linewidth=1.5, 
                label=label, zorder=0
                )
            # axes[j, i].text(
            #     0.01,
            #     thresh + 0.02,
            #     label,
            #     color=mcc_colors[label],
            #     ha="left",
            #     va="bottom",
            #     fontsize=10,
            #     fontweight="bold",
            #     transform=axes[j, i].transAxes,
            # )
        # If top row set the title
        if j == 0:
            axes[j, i].set_title(f"{pert}", fontsize=14, fontweight="bold")
        # For the last column, set the text for the dataType on the right
        if i == 3:
            axes[j, i].text(
                1.05,
                0.5,
                f"{dataType.capitalize()}",
                transform=axes[j, i].transAxes,
                ha="left",
                va="center",
                fontsize=14,
                fontweight="bold",
                rotation=270,
            )
# Add a text indicating the circle is the best MCC point under fig.title and above the first subplot
fig.text(
    0.9,
    .97,
    f"Black circles represent the maximum MCC values per method",
    ha="right",
    va="center",
    fontsize=12,
    fontstyle="italic",
    color="gray",
)
# Finalize the plot
plt.tight_layout()
plots.finalize_plot(
    fig, 
    show=True,
    save=save_to_folder,
    filename=f"{simID}_PepIDBenchmark_MCC_Curve", 
    filepath=figure_path,
    formats=figure_formats, 
    transparent=transparent_bg,
    dpi=figure_dpi,
)

# # MCC Score for different perturbations and methods at the best threshold

# Create subplots for MCC Score for different perturbations and methods at the best threshold
fig, axes = plt.subplots(
    1, 2, figsize=(12, 5), 
    sharey=True, sharex=True,
    gridspec_kw={"wspace": 0.05, "hspace": 0.1}
)

for i, dataType in enumerate(dataTypes):
    cur_data = idBenchmarkData[idBenchmarkData["DataType"] == dataType]
    sns.barplot(
        ax=axes[i],
        data=cur_data,
        x="Experiment",
        y="MCC",
        hue="Method",
        edgecolor="black",
        linewidth=0.8,
        errorbar="ci",  # Show confidence intervals across all thresholds
        capsize=0.05,
        err_kws={'linewidth': 1.5}
    )
    # Write the average MCC score for each method on the bottom of each bar
    for j, pert in enumerate(idBenchmarkData['Experiment'].unique()):
        for k, method in enumerate(method_palette.keys()):  # Changed from hardcoded list
            cur_score = cur_data[
                (cur_data["Experiment"] == pert) & 
                (cur_data["Method"] == method)
            ]["MCC"].mean()
            axes[i].text(
                j + k * 0.25 - 0.25,
                0 - 0.05,
                f"{cur_score:.2f}",
                color=method_palette[method],
                ha="center",
                va="top",
                fontsize=10,
                fontweight="bold",
                rotation=90,
            )


    # Get data specifically at the p-value threshold for point overlay
    pthr_data = cur_data[cur_data['threshold'] == pthr]
    # Add points for specific p-value threshold
    for j, method in enumerate(method_palette.keys()):
        method_data = pthr_data[pthr_data['Method'] == method]
        
        # Get x positions for this method's bars
        xlist = sorted(cur_data['Experiment'].unique())
        x_positions = []
        y_positions = []
        
        for k, x in enumerate(experiment_mapper.values()):
            mcc_at_pthr = method_data.loc[method_data['Experiment'] == x, 'MCC']
            if not mcc_at_pthr.empty:
                # Calculate x position (center of bar for this method)
                x_pos = k + (j - 1) * 0.27  # Adjust based on seaborn's bar positioning
                x_positions.append(x_pos)
                y_positions.append(mcc_at_pthr.iloc[0]) 
        
        # Add scatter points for this method at pthr
        axes[i].scatter(
            x_positions, 
            y_positions,
            color=method_palette[method],
            s=80,
            edgecolor='black',
            linewidth=1.25,
            marker='o',
            zorder=10,
            alpha=0.9
        )

    # Draw MCC interpretation thresholds
    for thresh, label in mcc_thresholds.items():
        axes[i].axhline(
            thresh, color=mcc_colors[label], alpha=1,
            linestyle="dotted", linewidth=1.5, 
            label=label, zorder=0
            )
            
    axes[i].set_title(f"{dataType.capitalize()}", fontsize=14, fontweight="bold")
    axes[i].set_xlabel('')
    axes[i].set_ylim(-0.15, 1.0)
    axes[i].grid("y", linestyle="--", linewidth=0.75, alpha=0.5, color="lightgrey")
    axes[i].set_xticklabels(axes[i].get_xticklabels(), rotation=0, ha="center", fontsize=10)
    # If first subplot, set the ylabel
    if i == 0:
        axes[i].set_ylabel("MCC Score", fontsize=12)
        # Remove the legend for the first subplot
        axes[i].legend().set_visible(False)
    else:
        axes[i].set_ylabel("")
        axes[i].legend(
            loc='upper right', title="Method", fontsize=10, title_fontsize=12, 
            frameon=False, 
        )

# Add a text to clarify the various annotations of the plot
# Include the p-value threshold in the text
# Include the MCC interpretation color legend
fig.text(
    0.9,
    .97,
    f"Black circles represent the MCC values at p-value threshold of {pthr}\nMCC interpretation thresholds are indicated by dotted lines",
    ha="right",
    va="center",
    fontsize=12,
    fontstyle="italic",
    color="gray",
)

plt.tight_layout()

plots.finalize_plot(
    fig, 
    show=True,
    save=save_to_folder,
    filename=f"{simID}_PepIDBenchmark_MCC_Barplot", 
    filepath=figure_path,
    formats=figure_formats, 
    transparent=transparent_bg,
    dpi=figure_dpi,
)

stTime = utils.getTime()

simID = "Sim2"  
# Set up a path for the simulation
output_path, figure_path = setup_simulation_paths( simID )

methods = list(method_palette.keys())
missRatePros = [0, 0.2, 0.4, 0.6, 0.8]
missRatePeps = [0, 0.2, 0.4, 0.6, 0.8]
missRate_mapper = {
    0: '0%',
    0.2: '20%',
    0.4: '40%',
    0.6: '60%',
    0.8: '80%',
}

print("PEPTIDE IDENTIFICATION BENCHMARK ANALYSIS")
print("-" * 52)
print(f"Simulation ID: {simID}")
print(f"Methods: {', '.join(methods)} ({len(methods)} total)")
print(f"Protein Missing Rates: {', '.join(map(str, missRatePros))}")
print(f"Peptide Missing Rates: {', '.join(map(str, missRatePeps))}")
print(f"Output Path: {output_path}")

# Progress tracking
total_combinations = len(methods) * len(missRatePros) * len(missRatePeps)
current_combination = 0
files_found = 0
files_missing = 0
print(f"Total combinations to process: {total_combinations}")
print("-" * 52)

display("Calculating ID Benchmarks:")

combined_results = []
for i in missRatePros:
    for j in missRatePeps:
        for cur_method in methods:
            current_combination += 1
            # Load the result data
            result_file = f"{output_path}2_Pro{i}_Pep{j}_imputed_{cur_method}_ResultData.feather"

            if os.path.exists(result_file):
                res_df = pd.read_feather(result_file)
                files_found += 1

                if verbose:
                    progress = f"[{current_combination:2d}/{total_combinations}]"
                    print(f"{progress} ✓ {cur_method:12} | ProMissing:{i:.1f} | PepMissing:{j:.1f}")
            else:
                files_missing += 1
                if verbose:
                    progress = f"[{current_combination:2d}/{total_combinations}]"
                    print(f"{progress} ✗ {cur_method:12} | ProMissing:{i:.1f} | PepMissing:{j:.1f} (MISSING)")
                continue
            
            # Handle COPF-specific column renaming
            if cur_method == "COPF":
                res_df = res_df.rename(columns={
                    "proteoform_score_pval": "adj_pval",
                    'protein_id': "Protein"
                })

            # Create metric data
            metric_data = utils.create_metric_data(
                res_df, 
                pvalue_thresholds=thresholds,
                label_col='pertPeptide',
                pvalue_col='adj_pval'
            )
            metric_data['Method'] = cur_method
            metric_data['ProteinMissingness'] = i
            metric_data['PeptideMissingness'] = j

            combined_results.append(metric_data)
    
# Combine and process data
idBenchmarkData = pd.concat(combined_results, ignore_index=True)
# Save the processed data
output_file = f"{output_path}4_{simID}_PeptideIdentification_PerformanceData.feather"
idBenchmarkData.to_feather(output_file)
idBenchmarkData.to_csv(output_file.replace('.feather', '.csv'), index=False)

print(f"\nDATA PREVIEW:")
display(idBenchmarkData.tail(5))

print(f"\nRESULTS SUMMARY:")
print("=" * 52)
execTime = utils.prettyTimer(utils.getTime() - stTime)
print(f"Total execution time:         {execTime}")
print(f"Files processed successfully: {files_found}")
print(f"Files missing/skipped:        {files_missing}")
print(f"Processing success rate:      {files_found/(files_found+files_missing)*100:.1f}%")
print("-" * 52)
print(f"Final dataset shape:          {idBenchmarkData.shape}")
protein_levels = sorted(idBenchmarkData['ProteinMissingness'].unique())
peptide_levels = sorted(idBenchmarkData['PeptideMissingness'].unique())
methods_list = sorted(idBenchmarkData['Method'].unique())

print(f"Unique protein missingness levels: {len(protein_levels)} -> {protein_levels}")
print(f"Unique peptide missingness levels: {len(peptide_levels)} -> {peptide_levels}")
print(f"Methods analyzed: {idBenchmarkData['Method'].nunique()} -> {', '.join(methods_list)}")
print(f"Threshold levels included: {idBenchmarkData['threshold'].nunique()} (min={idBenchmarkData['threshold'].min()}, max={idBenchmarkData['threshold'].max()})")
print(f"Data saved to: {output_file}")
print("=" * 52)

# Comparing AUC Values across different missingness levels
# Calculate AUC for ROC and PR curves
auc_results = []
for (method, pro_miss, pep_miss), group in idBenchmarkData.groupby(
        ['Method', 'ProteinMissingness', 'PeptideMissingness']
    ):
    # ROC AUC
    group = utils.complete_curve_data(group, 'ROC', 'FPR', 'TPR')
    roc_data = group.sort_values('FPR')
    roc_auc = np.trapezoid(roc_data['TPR'], roc_data['FPR'])
    
    
    auc_results.append({
        'Method': method,
        'ProteinMissingness': pro_miss,
        'PeptideMissingness': pep_miss,
        'AUC': roc_auc,
    })
auc_df = pd.DataFrame(auc_results)
print("\nAUC Data Preview:")
display(auc_df.head())

# Initialize Figure for 3 Methods with dedicated colorbar column
fig, axes = plt.subplots(
    nrows=1, ncols=4, figsize=(12, 4), 
    gridspec_kw={
        "wspace": 0.05, "hspace": 0.1,
        "width_ratios": [0.9, 0.9, 0.9, 0.1]  # First 3 columns same size, 4th for colorbar
    },
)
vmin = auc_df["AUC"].min()
vmax = auc_df["AUC"].max()

# Create heatmaps for the first 3 methods
for i, cur_method in enumerate(methods):
    plot_data = auc_df[auc_df["Method"] == cur_method].pivot_table(
        index="PeptideMissingness",
        columns="ProteinMissingness",
        values="AUC",
        aggfunc="mean"
    )
    # Make sure 0,0 is in the bottom left corner
    plot_data = plot_data.iloc[::-1]
    
    # Only show colorbar on the last method heatmap
    cb = (i == len(methods) - 1)
    
    # If this is the last method, create colorbar in the 4th column
    if cb:
        cbar_ax = axes[3]
    else:
        cbar_ax = None
    
    sns.heatmap(
        plot_data,
        vmin=vmin,
        vmax=vmax,
        cmap="Blues",
        annot=True,
        fmt=".2f",
        ax=axes[i],
        cbar=cb,
        cbar_ax=cbar_ax,
        cbar_kws={"label": "AUROC"},
        square=True,
        # Annotation size
        annot_kws={"size": 12}, 
    )
    
    if i == 0:
        axes[i].set_ylabel("Peptide Missingness Rate", fontsize=14, fontweight="bold")
    else:
        axes[i].set_ylabel("")
        # Remove yticklabels
        axes[i].set_yticklabels([])
    if i == 1:
        axes[i].set_xlabel("Protein Missingness Rate", fontsize=14, fontweight="bold")
    else:
        axes[i].set_xlabel("")

    axes[i].set_title(f"{cur_method}", fontsize=14, fontweight="bold")

# Hide the 4th axis if it wasn't used for colorbar (this handles the case where we have < 3 methods)
if len(methods) < 3:
    axes[3].set_visible(False)

plt.tight_layout()
plt.suptitle("Peptide Identification Benchmark with AUROC", fontsize=16, fontweight="bold")
plt.subplots_adjust(top=0.85)   

plots.finalize_plot(
    fig, 
    show=True,
    save=save_to_folder,
    filename=f"{simID}_PepIDBenchmark_AUROC_Heatmap", 
    filepath=figure_path,
    formats=figure_formats, 
    transparent=transparent_bg,
    dpi=figure_dpi,
)

# Mean MCC Scores
# Heatmap of the MCC values for the different methods 
# (Protein Missingness and Peptide Missingness) (X: Protein Missingness, Y: Peptide Missingness)

# Initalize Figure for 3 Methods
fig, axes = plt.subplots(
    nrows=1, ncols=4, figsize=(12, 4), 
    gridspec_kw={
        "wspace": 0.05, "hspace": 0.1,
        "width_ratios": [0.9, 0.9, 0.9, 0.1]  # First 3 columns same size, 4th for colorbar
    },
)

vmin = idBenchmarkData["MCC"].min()
vmax = idBenchmarkData["MCC"].max()

# Create heatmaps for the first 3 methods
for i, cur_method in enumerate(methods):
    plot_data = idBenchmarkData[idBenchmarkData["Method"] == cur_method].pivot_table(
        index="PeptideMissingness",
        columns="ProteinMissingness",
        values="MCC",
        aggfunc="mean"
    )
    # Make sure 0,0 is in the bottom left corner
    plot_data = plot_data.iloc[::-1]
    
    # Only show colorbar on the last method heatmap
    cb = (i == len(methods) - 1)
    
    # If this is the last method, create colorbar in the 4th column
    if cb:
        cbar_ax = axes[3]
    else:
        cbar_ax = None
    
    sns.heatmap(
        plot_data,
        vmin=vmin,
        vmax=vmax,
        cmap="Purples",
        annot=True,
        fmt=".2f",
        ax=axes[i],
        cbar=cb,
        cbar_ax=cbar_ax,
        cbar_kws={"label": "Mean MCC"},
        square=True,
        # Annotation size
        annot_kws={"size": 12}, 
    )
    
    if i == 0:
        axes[i].set_ylabel("Peptide Missingness Rate", fontsize=14, fontweight="bold")
    else:
        axes[i].set_ylabel("")
        # Remove yticklabels
        axes[i].set_yticklabels([])
    if i == 1:
        axes[i].set_xlabel("Protein Missingness Rate", fontsize=14, fontweight="bold")
    else:
        axes[i].set_xlabel("")

    axes[i].set_title(f"{cur_method}", fontsize=14, fontweight="bold")


# Hide the 4th axis if it wasn't used for colorbar (this handles the case where we have < 3 methods)
if len(methods) < 3:
    axes[3].set_visible(False)

plt.tight_layout()
plt.suptitle("Peptide Identification Benchmark with MCC (Mean)", fontsize=16, fontweight="bold")
plt.subplots_adjust(top=0.85)

plots.finalize_plot(
    fig, 
    show=True,
    save=save_to_folder,
    filename=f"{simID}_PepIDBenchmark_MCC_mean_Heatmap", 
    filepath=figure_path,
    formats=figure_formats, 
    transparent=transparent_bg,
    dpi=figure_dpi,
)

# MCC Score at pThr
# Heatmap of the MCC values for the different methods 
# (Protein Missingness and Peptide Missingness) (X: Protein Missingness, Y: Peptide Missingness)

# Initalize Figure for 3 Methods
fig, axes = plt.subplots(
    nrows=1, ncols=4, figsize=(12, 4), 
    gridspec_kw={
        "wspace": 0.05, "hspace": 0.1,
        "width_ratios": [0.9, 0.9, 0.9, 0.1]  # First 3 columns same size, 4th for colorbar
    },
)
vmin = idBenchmarkData["MCC"].min()
vmax = idBenchmarkData["MCC"].max()

# Create heatmaps for the first 3 methods
for i, cur_method in enumerate(methods):
    plot_data = idBenchmarkData[
        (idBenchmarkData["Method"] == cur_method) & 
        (idBenchmarkData["threshold"] == pthr)
    ].pivot_table(
        index="PeptideMissingness",
        columns="ProteinMissingness",
        values="MCC",
        aggfunc="mean"
    )
    # Make sure 0,0 is in the bottom left corner
    plot_data = plot_data.iloc[::-1]
    
    # Only show colorbar on the last method heatmap
    cb = (i == len(methods) - 1)
    
    # If this is the last method, create colorbar in the 4th column
    if cb:
        cbar_ax = axes[3]
    else:
        cbar_ax = None
    
    sns.heatmap(
        plot_data,
        vmin=vmin,
        vmax=vmax,
        cmap="Purples",
        annot=True,
        fmt=".2f",
        ax=axes[i],
        cbar=cb,
        cbar_ax=cbar_ax,
        cbar_kws={"label": f"MCC at p-value {pthr}"},
        square=True,
        # Annotation size
        annot_kws={"size": 12}, 
    )
    
    if i == 0:
        axes[i].set_ylabel("Peptide Missingness Rate", fontsize=14, fontweight="bold")
    else:
        axes[i].set_ylabel("")
        # Remove yticklabels
        axes[i].set_yticklabels([])
    if i == 1:
        axes[i].set_xlabel("Protein Missingness Rate", fontsize=14, fontweight="bold")
    else:
        axes[i].set_xlabel("")

    axes[i].set_title(f"{cur_method}", fontsize=14, fontweight="bold")

# Hide the 4th axis if it wasn't used for colorbar (this handles the case where we have < 3 methods)
if len(methods) < 3:
    axes[3].set_visible(False)  

plt.tight_layout()
plt.suptitle(f"Peptide Identification Benchmark with MCC at p-value {pthr}", fontsize=16, fontweight="bold")
plt.subplots_adjust(top=0.85)   

plots.finalize_plot(
    fig, 
    show=True,
    save=save_to_folder,
    filename=f"{simID}_PepIDBenchmark_MCC_pThr_Heatmap", 
    filepath=figure_path,
    formats=figure_formats, 
    transparent=transparent_bg,
    dpi=figure_dpi,
)

# Maximum MCC Scores
# Heatmap of the MCC values for the different methods 
# (Protein Missingness and Peptide Missingness) (X: Protein Missingness, Y: Peptide Missingness)

# Initialize Figure for 3 Methods with dedicated colorbar column
fig, axes = plt.subplots(
    nrows=1, ncols=4, figsize=(12, 4), 
    gridspec_kw={
        "wspace": 0.05, "hspace": 0.1,
        "width_ratios": [0.9, 0.9, 0.9, 0.1]  # First 3 columns same size, 4th for colorbar
    },
)
vmin = idBenchmarkData["MCC"].min()
vmax = idBenchmarkData["MCC"].max()

# Create heatmaps for the first 3 methods
for i, cur_method in enumerate(methods):
    plot_data = idBenchmarkData[idBenchmarkData["Method"] == cur_method].pivot_table(
        index="PeptideMissingness",
        columns="ProteinMissingness",
        values="MCC",
        aggfunc="max"
    )
    # Make sure 0,0 is in the bottom left corner
    plot_data = plot_data.iloc[::-1]
    
    # Only show colorbar on the last method heatmap
    cb = (i == len(methods) - 1)
    
    # If this is the last method, create colorbar in the 4th column
    if cb:
        cbar_ax = axes[3]
    else:
        cbar_ax = None
    
    sns.heatmap(
        plot_data,
        vmin=vmin,
        vmax=vmax,
        cmap="Purples",
        annot=True,
        fmt=".2f",
        ax=axes[i],
        cbar=cb,
        cbar_ax=cbar_ax,
        cbar_kws={"label": "Max MCC"},
        square=True,
        # Annotation size
        annot_kws={"size": 12}, 
    )
    
    if i == 0:
        axes[i].set_ylabel("Peptide Missingness Rate", fontsize=14, fontweight="bold")
    else:
        axes[i].set_ylabel("")
        # Remove yticklabels
        axes[i].set_yticklabels([])
    if i == 1:
        axes[i].set_xlabel("Protein Missingness Rate", fontsize=14, fontweight="bold")
    else:
        axes[i].set_xlabel("")

    axes[i].set_title(f"{cur_method}", fontsize=14, fontweight="bold")


# Hide the 4th axis if it wasn't used for colorbar (this handles the case where we have < 3 methods)
if len(methods) < 3:
    axes[3].set_visible(False)

plt.tight_layout()
plt.suptitle("Peptide Identification Benchmark with MCC (Max)", fontsize=16, fontweight="bold")
plt.subplots_adjust(top=0.85)

plots.finalize_plot(
    fig, 
    show=True,
    save=save_to_folder,
    filename=f"{simID}_PepIDBenchmark_MCC_max_Heatmap", 
    filepath=figure_path,
    formats=figure_formats, 
    transparent=transparent_bg,
    dpi=figure_dpi,
)

stTime = utils.getTime()

simID = "Sim3"  
# Set up a path for the simulation
output_path, figure_path = setup_simulation_paths( simID )

methods = list(method_palette.keys())
perturbationRanges = [
    (0.1, 0.25),
    (0.25, 0.50),
    (0.50, 0.75),
    (0.75, 1.0),
    (1.0, 1.25),
    (1.25, 1.50),
    (1.50, 1.75),
    (1.75, 2.0),
]

perturbation_mapper = {
    (0.1, 0.25): '0.1-0.25',
    (0.25, 0.50): '0.25-0.50',
    (0.50, 0.75): '0.50-0.75',
    (0.75, 1.0): '0.75-1.0',
    (1.0, 1.25): '1.0-1.25',
    (1.25, 1.50): '1.25-1.50',
    (1.50, 1.75): '1.50-1.75',
    (1.75, 2.0): '1.75-2.0',
}


print("PEPTIDE IDENTIFICATION BENCHMARK ANALYSIS")
print("-" * 52)
print(f"Simulation ID: {simID}")
print(f"Methods: {', '.join(methods)} ({len(methods)} total)")
print(f"Perturbation Ranges: {', '.join([f'{low}-{high}' for low, high in perturbationRanges])}")
print(f"Output Path: {output_path}")

# Progress tracking
total_combinations = len(methods) * len(perturbationRanges)
current_combination = 0
files_found = 0
files_missing = 0
print(f"Total combinations to process: {total_combinations}")
print("-" * 52)
display("Calculating ID Benchmarks:")

combined_results = []
for low, high in perturbationRanges:
    for cur_method in methods:
        current_combination += 1
        # Load the result data
        result_file = f"{output_path}2_{low}_{high}_{cur_method}_ResultData.feather"

        if os.path.exists(result_file):
            res_df = pd.read_feather(result_file)
            files_found += 1

            if verbose:
                progress = f"[{current_combination:2d}/{total_combinations}]"
                print(f"{progress} ✓ {cur_method:12} | Perturbation: {low:.2f}-{high:.2f}")
        else:
            files_missing += 1
            if verbose:
                progress = f"[{current_combination:2d}/{total_combinations}]"
                print(f"{progress} ✗ {cur_method:12} | Perturbation: {low:.2f}-{high:.2f} (MISSING)")
            continue
        
        # Handle COPF-specific column renaming
        if cur_method == "COPF":
            res_df = res_df.rename(columns={
                "proteoform_score_pval": "adj_pval",
                'protein_id': "Protein"
            })

        # Create metric data
        metric_data = utils.create_metric_data(
            res_df, 
            pvalue_thresholds=thresholds,
            label_col='pertPeptide',
            pvalue_col='adj_pval'
        )
        metric_data['Method'] = cur_method
        metric_data['PerturbationRange'] = f"{low:.2f}-{high:.2f}"

        combined_results.append(metric_data)

# Combine and process data
idBenchmarkData = pd.concat(combined_results, ignore_index=True)
# Save the processed data
output_file = f"{output_path}4_{simID}_PeptideIdentification_PerformanceData.feather"
idBenchmarkData.to_feather(output_file) 
idBenchmarkData.to_csv(output_file.replace('.feather', '.csv'), index=False)

print(f"\nDATA PREVIEW:")
display(idBenchmarkData.tail(5))

print(f"\nRESULTS SUMMARY:")
print("=" * 52)
execTime = utils.prettyTimer(utils.getTime() - stTime)
print(f"Total execution time:         {execTime}")
print(f"Files processed successfully: {files_found}")
print(f"Files missing/skipped:        {files_missing}")
print(f"Processing success rate:      {files_found/(files_found+files_missing)*100:.1f}%")
print("-" * 52)
print(f"Final dataset shape:          {idBenchmarkData.shape}")
pert_ranges = sorted(idBenchmarkData['PerturbationRange'].unique())
methods_list = sorted(idBenchmarkData['Method'].unique())
print(f"Unique perturbation ranges: {len(pert_ranges)} -> {', '.join(pert_ranges)}")
print(f"Methods analyzed: {idBenchmarkData['Method'].nunique()} -> {', '.join(methods_list)}")
print(f"Threshold levels included: {idBenchmarkData['threshold'].nunique()} (min={idBenchmarkData['threshold'].min()}, max={idBenchmarkData['threshold'].max()})")
print(f"Data saved to: {output_file}")
print("=" * 52)

# Line plot showing MCC scores across perturbation ranges 
# line plot marker at the mean MCC and CI as a shaded area
# for each method and data type
fig, ax = plt.subplots(1, 1, figsize=(12, 6))

sns.lineplot(
    ax=ax,
    data=idBenchmarkData,
    x="PerturbationRange",
    y="MCC",
    hue="Method",
    style="Method",
    palette=method_palette,
    linewidth=3,
    markers=method_markers,
    markersize=12,
    markeredgewidth=0.5,
    markeredgecolor="black", 
    
    # Estimation and confidence interval
    estimator='mean',
    errorbar=('ci', 95),
    err_style="band",
    rasterized=True,
)

ax.set_xlabel("Perturbation Range", fontsize=14, fontweight="bold")
ax.set_ylabel("MCC Score", fontsize=14, fontweight="bold")
ax.set_title("Peptide Identification Benchmark with MCC across Perturbation Ranges", fontsize=16, fontweight="bold")
ax.set_ylim(-0.15, 1.0)
ax.grid("y", linestyle="--", linewidth=0.75, alpha=0.5, color="lightgrey")
# ax.set_xticklabels(ax.get_xticklabels(), rotation=45, ha="right", fontsize=10)
ax.legend(
    loc='upper right', title="Method", fontsize=10, title_fontsize=12, 
    frameon=False, 
)   

# Draw MCC interpretation thresholds
for thresh, label in mcc_thresholds.items():
    ax.axhline(
        thresh, color=mcc_colors[label], alpha=1,
        linestyle="dotted", linewidth=1.5, 
        label=label, zorder=0
        )
    # Use mixed transform: relative x, data y
    ax.text(
        0.01,
        thresh,
        label,
        color=mcc_colors[label],
        ha="left",
        va="center",  # Center on the line
        fontsize=10,
        fontweight="bold",
        transform=ax.get_yaxis_transform(),  # x in axes coords, y in data coords
        bbox=dict(boxstyle='round,pad=0.3', facecolor='white', edgecolor=mcc_colors[label], alpha=0.8)
    )

# Add smaller scatter points for MCC values at p-value threshold
pthr_data = idBenchmarkData[idBenchmarkData['threshold'] == pthr]
for i, method in enumerate(method_palette.keys()):
    method_data = pthr_data[pthr_data['Method'] == method]
    ax.scatter(
        method_data['PerturbationRange'],
        method_data['MCC'],
        color=method_palette[method],
        s=50,
        edgecolor='black',
        linewidth=1.25,
        marker=method_markers[method],
        zorder=10,
        alpha=0.9
    )

# Add a text to clarify the various annotations of the plot
# Include the p-value threshold in the text
fig.text(
    1,
    1.05,
    f"Black markers represent the MCC values at p-value threshold of {pthr}\nMCC interpretation thresholds are indicated by dotted lines\nShaded areas represent 95% confidence intervals around the mean MCC",
    ha="right",
    va="center",
    fontsize=12,
    fontstyle="italic",
    color="gray",
)
# Finalize the plot
plt.tight_layout()
plots.finalize_plot(
    fig, 
    show=True,
    save=save_to_folder,
    filename=f"{simID}_PepIDBenchmark_MCC_Curve_Perturbation", 
    filepath=figure_path,
    formats=figure_formats, 
    transparent=transparent_bg,
    dpi=figure_dpi,
)